// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <cstdint>

#include "melon/io/cursor.h"
#include <pollux/common/file/file.h>
#include <pollux/common/file/region.h>

namespace kumo::pollux::file::utils {
    // Iterable class that produces pairs of iterators pointing to the beginning and
    // end of the segments that are coalesced from the the input range, according to
    // the ShouldCoalesce condition
    template<typename RegionIter, typename ShouldCoalesce>
    class CoalesceRegions {
    public:
        class Iter {
        public:
            Iter(RegionIter begin, RegionIter end, ShouldCoalesce shouldCoalesce)
                : begin_{begin},
                  end_{end},
                  theEnd_{end},
                  shouldCoalesce_(std::move(shouldCoalesce)) {
                findNextEnd();
            }

            friend bool operator==(const Iter &lhs, const Iter &rhs) {
                return lhs.begin_ == rhs.begin_ && lhs.end_ == rhs.end_;
            }

            friend bool operator!=(const Iter &lhs, const Iter &rhs) {
                return !(lhs == rhs);
            }

            std::pair<RegionIter, RegionIter> operator*() const {
                return {begin_, end_};
            }

            Iter operator++() {
                begin_ = end_;
                end_ = theEnd_;
                findNextEnd();
                return *this;
            }

            Iter operator++(int) {
                Iter tmp(*this);
                ++(*this);
                return tmp;
            }

        private:
            void findNextEnd() {
                if (begin_ != theEnd_) {
                    for (auto itA = begin_, itB = std::next(itA); itB != theEnd_;
                         itA = itB, ++itB) {
                        if (!shouldCoalesce_(*itA, *itB)) {
                            end_ = itB;
                            break;
                        }
                    }
                }
            }

            RegionIter begin_;
            RegionIter end_;
            RegionIter theEnd_;
            ShouldCoalesce shouldCoalesce_;
        };

        CoalesceRegions(
            RegionIter begin,
            RegionIter end,
            ShouldCoalesce shouldCoalesce)
            : begin_{begin}, end_{end}, shouldCoalesce_(std::move(shouldCoalesce)) {
        }

        Iter begin() {
            return Iter{begin_, end_, shouldCoalesce_};
        }

        Iter end() {
            return Iter{end_, end_, shouldCoalesce_};
        }

    private:
        RegionIter begin_;
        RegionIter end_;
        ShouldCoalesce shouldCoalesce_;
    };

    class CoalesceIfDistanceLE {
    public:
        explicit CoalesceIfDistanceLE(
            uint64_t maxCoalescingDistance,
            uint64_t * MELON_NULLABLE coalescedBytes = nullptr)
            : maxCoalescingDistance_{maxCoalescingDistance},
              coalescedBytes_{coalescedBytes} {
        }

        bool operator()(
            const pollux::common::Region &a,
            const pollux::common::Region &b) const;

    private:
        uint64_t maxCoalescingDistance_;
        uint64_t *coalescedBytes_;
    };

    template<typename RegionIter, typename OutputIter, typename Reader>
    class ReadToIOBufs {
    public:
        ReadToIOBufs(
            RegionIter begin,
            RegionIter end,
            OutputIter output,
            Reader reader)
            : begin_{begin}, end_{end}, output_{output}, reader_{std::move(reader)} {
        }

        void operator()() {
            if (begin_ == end_) {
                return;
            }

            auto fileOffset = begin_->offset;
            const auto last = std::prev(end_);
            const auto readSize = last->offset + last->length - fileOffset;
            std::unique_ptr<melon::IOBuf> result = reader_(fileOffset, readSize);

            melon::io::Cursor cursor(result.get());
            for (auto region = begin_; region != end_; ++region) {
                if (fileOffset < region->offset) {
                    cursor.skip(region->offset - fileOffset);
                    fileOffset = region->offset;
                }
                // This clone won't copy the underlying buffer. It will just create an
                // IOBuf pointing to the right section of the existing shared buffer
                // of the original IOBuf. It can create a chained IOBuf if the original
                // IOBuf is chained, and the length of the current read spreads to the
                // next IOBuf in the chain.
                melon::IOBuf buf;
                cursor.clone(buf, region->length);
                *output_++ = std::move(buf);
                fileOffset += region->length;
            }
        }

    private:
        RegionIter begin_;
        RegionIter end_;
        OutputIter output_;
        Reader reader_;
    };
} // namespace kumo::pollux::file::utils
